import torch
import torch.nn as nn
import torch.fx
from parse_config import Config
import shutil
import os
import model_input


def create_graph(save_path):
    # 创建模型实例
    model = model_input.get_model()

    traced_model = torch.fx.symbolic_trace(model)

    # 保存计算图到文件
    torch.save(traced_model, save_path)



def main():
    shutil.rmtree('./results', ignore_errors=True)
    os.makedirs('./results',exist_ok=True)
    config_path='./config.json'
    config=Config(config_path)
    create_graph(config.config['formal_model_save_path'])
    
if __name__ == '__main__':
    main()
    



    